from mindspore import nn


class Schrodinger_Net(nn.Cell):
    def __init__(self):
        super(Schrodinger_Net, self).__init__()
        self.fc1 = nn.Dense(2, 100, activation='tanh', weight_init='xavier_uniform')
        self.fc2 = nn.Dense(100, 100, activation='tanh', weight_init='xavier_uniform')
        self.fc3 = nn.Dense(100, 100, activation='tanh', weight_init='xavier_uniform')
        self.fc4 = nn.Dense(100, 100, activation='tanh', weight_init='xavier_uniform')
        self.fc5 = nn.Dense(100, 2, weight_init='xavier_uniform')

    def construct(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.fc4(x)
        x = self.fc5(x)
        return x